import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt


def get_cifar_data(batch_size):
    # data compose
    transform = transforms.Compose([
        # padding to 36x36
        transforms.Pad(4),
        # Random Horizontal Flip
        transforms.RandomHorizontalFlip(),
        # cut to 32x32
        transforms.RandomCrop(32),
        # to tensor
        transforms.ToTensor(),
        # normalization
        transforms.Normalize(mean=(0.5, 0.5, 0.5),
                             std=(0.5, 0.5, 0.5))
    ])

    # cifar10 path
    cifar10Path = './cifar'

    # train data
    train_dataset = torchvision.datasets.CIFAR10(root=cifar10Path,
                                                 train=True,
                                                 transform=transform,
                                                 download=True)

    # test data
    test_dataset = torchvision.datasets.CIFAR10(root=cifar10Path,
                                                train=False,
                                                transform=transform)

    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,
                                               shuffle=True)

    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                              batch_size=batch_size,
                                              shuffle=False)

    return train_loader, test_loader


if __name__ == '__main__':
    train_loader, test_loader = get_cifar_data(batch_size=32)

    data_iter = iter(train_loader)
    images, labels = next(data_iter)
    idx = 31
    image = images[idx].numpy()
    image = np.transpose(image, (1, 2, 0))
    plt.imshow(image)
    plt.show()
    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    print(classes[labels[idx].numpy()])
    print("1")


    print(type(train_loader))